
# sudo sh -c 'echo "export OLLAMA_HOST=0.0.0.0:11434" >> /etc/profile'
# launchctl setenv OLLAMA_HOST "0.0.0.0:11434"

import requests
import json

# 定义 Ollama API 的基本 URL，假设 Ollama 在本地运行，端口为 11434
BASE_URL = "http://localhost:11434"

class LLM_Chat:
    def __init__(self, model: str = "qwen2.5:0.5b", host: str = BASE_URL):
        self.model = model
        self.host = f"{host}/api/chat"
       
    def chat(self, messages):
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.7,
            "stream": False  # 设置为 False 获取完整响应
        }

        try:
            # 发送 POST 请求
            response = requests.post(self.host, json=payload)
            response.raise_for_status()  # 如果请求失败 (状态码 >= 400), 抛出 HTTPError
            # 解析响应
            result = response.json()
            print(result)
            return result["message"]["content"]  
        except Exception as e:
            return f"发生未知错误: {e}"
        
    def test(self):
        # 示例：构建包含历史消息的对话
        history_messages = [
            {"role": "user", "content": "你能帮我解释一下机器学习的基本概念吗？"},
            {"role": "assistant", "content": "当然可以！机器学习是人工智能的一个分支，它通过数据和算法让计算机自动学习模式并做出预测。"},
            {"role": "user", "content": "监督学习和无监督学习有什么区别？"}
        ]

        # 调用 Ollama
        response = self.chat(
            messages=history_messages
        )
        
# --- 使用示例 ---
if __name__ == "__main__":
    llm = LLM_Chat()
    llm.test()
    
    
